import os
import numpy as np
import json
import open3d as o3d
from scipy.spatial.transform import Rotation as R

cur_dir = os.path.dirname(__file__)


with open("./BoneInfo.json", "r") as f:
    boneinfos = json.load(f)
root_index = 0

for frame_index in range(0, 100000, 10):
    bones, bone_poses = [], []
    with open("/home/liuyun/Humanoid_IL_Benchmark/Kinematic_Data/ROAM/chair_sit/HumanPoses1-2800/1a6f615e8b1b5ae4dbbc9440457e303e_pose34_coord_scale/Positions.json", "r") as f:
        total_position = json.load(f)
        print(frame_index)
        if frame_index >= len(total_position):
            break
        bone_position = total_position[frame_index]
        
    with open("/home/liuyun/Humanoid_IL_Benchmark/Kinematic_Data/ROAM/chair_sit/HumanPoses1-2800/1a6f615e8b1b5ae4dbbc9440457e303e_pose34_coord_scale/Rotations.json", "r") as f:
        bone_rotations = json.load(f)[frame_index]

        def GetTransformationByBoneIndex(index):
            transform = np.eye(4)
            pos = boneinfos[index]["InitPosition"] if index != root_index else bone_position[index]
            transform[0][3] = pos["x"]
            transform[1][3] = pos["y"]
            transform[2][3] = pos["z"]
            transform[:3, :3] = GetRotationByBoneIndex(index)
            return transform
        
        def GetRotationByBoneIndex(index):
            quat = np.array([bone_rotations[index]["x"], bone_rotations[index]["y"], bone_rotations[index]["z"], bone_rotations[index]["w"]])
            return R.from_quat(quat).as_matrix()

        
        def compute_global_transform(bone, parent_transform):
            local_transform = GetTransformationByBoneIndex(bone["Index"])
            global_transform = np.dot(parent_transform, local_transform)
            # print(bone)
            # print(parent_transform)
            # print(local_transform)
            # print(global_transform)
            # input()
            return global_transform

        # 定义一个函数来遍历关节树并计算关节位置
        def traverse_skeleton(bone, parent_transform, joint_positions, joint_poses):
            global_transform = compute_global_transform(bone, parent_transform)
            joint_positions[bone["Index"]] = global_transform[:, 3][:3]
            joint_poses[bone["Index"]] = global_transform.copy()
            for child in bone["Childs"]:
                traverse_skeleton(boneinfos[child], global_transform, joint_positions, joint_poses)

        final_pos = [np.zeros(3) for _ in range(len(boneinfos))]
        final_pose = [np.eye(4) for _ in range(len(boneinfos))]
        traverse_skeleton(boneinfos[root_index], np.eye(4), final_pos, final_pose)
        for pos in final_pos:
            bones.append(pos)
        for pose in final_pose:
            bone_poses.append(pose)
    new_bones = []
    for pos in bone_position:
        new_bones.append(np.array([pos["x"], pos["y"], pos["z"]]))
    # for i in range(len(bones)):
    #     print("{} / {} / {}".format(i, bones[i], new_bones[i]))

    coords = []
    for i in range(len(bones)):
        coord = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.05)
        coord.transform(bone_poses[i])
        coords.append(coord)

    # 将所有立方体合并为一个列表
    # all_cubes = o3d.geometry.TriangleMesh.create_concatenated_mesh(cubes)
    transform_matrix = np.load("/home/liuyun/Humanoid_IL_Benchmark/Kinematic_Data/ROAM/chair_sit/HumanPoses1-2800/1a6f615e8b1b5ae4dbbc9440457e303e_pose34_coord_scale/ObjectPose.npy")
    object = o3d.io.read_triangle_mesh("/home/liuyun/Humanoid_IL_Benchmark/Kinematic_Data/ROAM/chair_sit/ObjectDemoCollections1-2800/1a6f615e8b1b5ae4dbbc9440457e303e_pose34_coord_scale_mesh_world.obj").transform(transform_matrix)
    coords.append(object)
    # 设置Open3D窗口的背景颜色
    o3d.visualization.draw_geometries(coords)
